from pymilvus import MilvusClient, Collection, FieldSchema, CollectionSchema, DataType
from config import milvus_url, milvus_db_name, milvus_limit, ollama_host, ollama_port
from langchain_ollama import OllamaEmbeddings
client = MilvusClient(
    uri=milvus_url,
    db_name=milvus_db_name,
)

def creat_tools_info_vector(columns_collection):

    client.drop_collection(columns_collection)
    if client.has_collection(columns_collection):
        client.load_collection(columns_collection)
    else:
        fields = [
            FieldSchema(name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100),
            FieldSchema(name="tool_name", dtype=DataType.VARCHAR, max_length=32),
            FieldSchema(name="desc", dtype=DataType.VARCHAR, max_length=2000),
            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024)
        ]
        schema = CollectionSchema(fields, columns_collection)
        client.create_collection(columns_collection, schema=schema)
        index_params = client.prepare_index_params()
        index_params.add_index(
            field_name="vector",
            index_type="IVF_FLAT",
            metric_type="COSINE",
            params={
                "nprobe": 1024,  # 用于搜索的聚类单元数量
                "radius": 0.2,  # 设置最小相似度阈值 - range_filter: 用于在特定相似度范围内细化搜索
                "range_filter": 1.0,  # range_filter 的值应该大于 radius 的值
                "nlist":1024
            }
            # metric_type="L2",
            # params={"nlist": 128}
        )
        client.create_index(columns_collection, index_params=index_params)
        client.load_collection(columns_collection)

def add_document(collection_name, data):
    client.insert(collection_name, data)

def embedding(query):
    v = OllamaEmbeddings(
        base_url=ollama_host + ':' + str(ollama_port),
        model="bge-m3:latest",
    )

    return v.embed_query(query)


if __name__ == '__main__':
    collection_name = "tool_map_desc"
    creat_tools_info_vector(collection_name)
    # data = {
    #     "vector": embedding(question)
    # }
    l = [["get_upland_water_info",     """
    查询上游来水数据
    :param company_id:公司id
    :param start_date:
    :param end_date:
    :return:
    """],["get_upland_water_info",     """
    查询下游排海数据
    :param start_date:
    :param end_date:
    :return:
    """], ["get_lab_internal_info",     """
    根据起止日期获取实验室在线数据、检测分析数据，其中参数可以为空，为空参数设置为None
    :param checkStatus:实验室检测状态列表["正常","异常"], default:None
    :param sampleId:实验室检测位置列表["排放池", "磁混凝","臭氧池", "曝气生物滤池", "反硝化滤池(B)","反硝化滤池(A)","二期调节池","高效沉淀池","二沉池(B)","二沉池(A)","生化池1(ASR-B)","生化池1(ASR-A)","缺氧池(B)","缺氧池(A)","AnGSB池(B)","AnGSB池(A)","高浓度调节池","低浓度调节池","园区总管"], default:None
    :param start_date:
    :param end_date:
    :return:
    """], ["get_device_check",     """
    查询厂内运行巡检(巡检缺卡/运行缺卡)数据，其中参数可以为空，为空参数设置为None
    :param shift:班次列表["A班组","B班组", "C班组", "D班组"], default:None
    :param type1:["2H类", "12H类"], default:None
    :param appealStatus:["未申诉", "申述中","申述成功", "申述失败", "申述驳回"], default:None
    :param startTime:
    :param endTime:
    :return:
    """], ["device_check_absence",    """
    查询厂内运行巡检(巡检打卡/运行打卡、巡检正常/异常情况、巡检异常维修情况)数据
    :param repairState:厂内修理状态列表["无需维修","维修单生成", "维修已完成", "正在维修"], default:None
    :param abnormal:厂内是否异常列表["正常情况", "异常情况"], default:None
    :param startTime:
    :param endTime:
    :return:
    """],["get_internal_online_data",     """
    查询厂内在线数据
    :param startTime:
    :param endTime:
    :return:
    """],["get_electric_consume_record",    """
    查询电量能耗数据
    :param startTime:
    :param endTime:
    :param type1:设备类型列表["1#机主","2#机主", "深度处理A", "深度处理B", "新风机A", "新风机B", "新风机C"], default:None
    :return:
    """ ], ["get_composite_consume_record",     """
    查询用水量(工业用水、生活用水)能耗数据
    :param startTime:
    :param endTime:
    :param type1:用水类型列表["工业用水","生活用水"], default:None
    :return:
    """],["get_steam_dosage_record",     """
    查询用蒸汽量能耗数据
    :param startTime:
    :param endTime:
    :return:
    """], ["get_drug_inventory",     """
    查询药剂管控(出库、入库、库存、供应商)数据
    :param startTime:
    :param endTime:
    :param name: 药品名称列表['液体柠檬酸', '次钠', '高效复合碳源', '乙酸钠', '次氯酸钠固体', '消泡剂', '液体尿素', '尿素', '磷酸氢二铵', 'PAM(S806)(阴)', 'PAM(4490)(阳)', 'PAC', '液碱', '液氧', '活性炭', '营养液', '硫酸亚铁'] default:None
    :param type1:进出库状态列表['出库','入库', '出入库'], default:出入库
    :param deliverer:公司名称列表['信安贸易有限公司', '嘉兴市嘉赢环境工程有限公司', '浙江巨能环境工程有限公司', '嘉兴市浩晖贸易有限公司', '嘉兴市南湖区首创助剂厂', '上海协聪化工有限公司', '嘉兴市堃达化工有限公司', '林德（上海）工业气体有限公司', '浙江嘉通新材料科技有限公司', '江苏正力环保科技有限公司', '南京欣瑞新材料科技有限公司', '嘉兴市禾晟生物制品有限公司', '嘉兴市创伟新材料股份有限公司']default:None
    :return:
    """], ["get_sludge_transport_record",     """
    按按进场/入库时间查询污泥外运数据以及污泥含水量数据
    :param timeInStart:
    :param timeInEnd:
    :param destination:['嘉爱斯', '平湖石化'] default:None
    :return:
    """], ["get_sludge_transport_record_out",  """
    按出场/出库时间查询污泥外运数据以及污泥含水量外运数据
    :param timeOutStart:
    :param timeOutEnd:
    :param destination:['嘉爱斯', '平湖石化'] default:None
    :return:
    """], ["get_sludge_transport_record",     """
    查询人员考勤打卡数据
    :param startTime:
    :param endTime:
    :param timeResult: 打卡情况列表['已打卡', '缺卡/未打卡', '早退', '迟到', '严重迟到', '旷工迟到', '休息'] default:None
    :return:
    """],["get_device",   """
    查询设备清单(运行中设备、维修中设备、异常状态设备、停用设备)/设备详情数据
    :param state: 设备状态列表['运行中', '维修中/故障情况', '异常', '停用']
    :return:
    """]]
    for a in l:
        data = {
            "tool_name": a[0],
            "desc":a[1],
            "vector": embedding(a[1])
        }
        add_document(collection_name, data)